package com.agent.platform.service;

import com.agent.platform.client.ContextManagerClient;
import com.agent.platform.client.ToolManagerClient;
import com.agent.platform.entity.Agent;
import com.agent.platform.entity.AgentSession;
import com.agent.platform.repository.AgentSessionRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 会话循环执行服务 - 独立Service以支持@Async
 *
 * 从AgentSessionService中分离出来，解决同类调用@Async不生效的问题
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class SessionLoopService {

    private final AgentSessionRepository sessionRepository;
    private final ModelSelectorService modelSelectorService;
    private final ContextManagerClient contextManagerClient;
    private final ToolManagerClient toolManagerClient;

    // 运行中的会话标记
    private final Map<String, Boolean> runningSessions = new ConcurrentHashMap<>();

    /**
     * 异步启动会话循环
     */
    @Async
    public void startSessionLoop(Agent agent, AgentSession session) {
        String sessionId = session.getSessionId();
        String agentId = agent.getAgentId();

        runningSessions.put(sessionId, true);

        log.info("Session loop started for: {}", sessionId);

        try {
            while (runningSessions.getOrDefault(sessionId, false)) {
                // 检查最大步数
                if (session.getStepCount() >= agent.getMaxSteps()) {
                    log.warn("Max steps reached for session: {}", sessionId);
                    break;
                }

                // 1. 获取新消息
                Optional<ContextManagerClient.MessageData> newMessage = fetchNewMessage(agent, session);

                if (newMessage.isEmpty()) {
                    // 没有新消息，等待后继续
                    Thread.sleep(2000);
                    continue;
                }

                // 2. 处理消息
                processMessage(agent, session, newMessage.get());

                // 3. 更新会话
                session = sessionRepository.save(session);

                // 检查失败次数
                if (session.getFailureCount() >= agent.getMaxFailures()) {
                    log.error("Max failures reached for session: {}", sessionId);
                    session.markAsFailed("Max failures reached");
                    sessionRepository.save(session);
                    break;
                }

                // 短暂休息
                Thread.sleep(500);
            }

        } catch (Exception e) {
            log.error("Session loop failed for: {}", sessionId, e);
            session.markAsFailed(e.getMessage());
            sessionRepository.save(session);
        } finally {
            runningSessions.remove(sessionId);
            log.info("Session loop ended for: {}", sessionId);
        }
    }

    /**
     * 停止指定会话的循环
     */
    public void stopSessionLoop(String sessionId) {
        runningSessions.remove(sessionId);
        log.info("Stopping session loop for: {}", sessionId);
    }

    /**
     * 检查会话是否正在运行
     */
    public boolean isSessionRunning(String sessionId) {
        return runningSessions.getOrDefault(sessionId, false);
    }

    /**
     * 获取新消息
     */
    private Optional<ContextManagerClient.MessageData> fetchNewMessage(Agent agent, AgentSession session) {
        try {
            ContextManagerClient.GetMessagesResponse response = contextManagerClient
                .getMessages(
                    agent.getContextManagerUrl(),
                    session.getSessionId(),
                    agent.getAgentId(),
                    10,
                    0
                )
                .block();

            if (response != null && response.getMessages() != null && !response.getMessages().isEmpty()) {
                // 返回最新的用户消息
                return response.getMessages().stream()
                    .filter(msg -> "user".equals(msg.getRole()))
                    .findFirst();
            }

        } catch (Exception e) {
            log.error("Failed to fetch messages", e);
        }

        return Optional.empty();
    }

    /**
     * 处理消息
     */
    private void processMessage(
        Agent agent,
        AgentSession session,
        ContextManagerClient.MessageData userMessage
    ) {
        try {
            session.incrementStepCount();

            // 1. 构建消息列表
            List<Message> messages = new ArrayList<>();
            messages.add(new UserMessage(userMessage.getContent()));

            // 2. 获取可用工具
            List<Map<String, Object>> tools = fetchTools(agent, session);

            // 3. 调用模型
            String model = modelSelectorService.selectModel(session.getMetadata());
            ChatResponse response = modelSelectorService.callModel(messages, tools, model);

            // 4. 处理响应
            String content = response.getResult().getOutput().getContent();

            // 5. 保存响应到上下文管理器
            ContextManagerClient.MessageData assistantMessage = new ContextManagerClient.MessageData();
            assistantMessage.setRole("assistant");
            assistantMessage.setContent(content);

            contextManagerClient.addMessage(
                agent.getContextManagerUrl(),
                session.getSessionId(),
                agent.getAgentId(),
                assistantMessage
            ).block();

            // 6. 更新统计
            if (response.getMetadata() != null && response.getMetadata().getUsage() != null) {
                session.addTokens(Math.toIntExact(response.getMetadata().getUsage().getTotalTokens()));
            }

            log.info("Processed message for session: {} step: {}", session.getSessionId(), session.getStepCount());

        } catch (Exception e) {
            log.error("Failed to process message", e);
            session.incrementFailureCount();
            session.setLastError(e.getMessage());
        }
    }

    /**
     * 获取工具列表
     */
    private List<Map<String, Object>> fetchTools(Agent agent, AgentSession session) {
        if (!agent.isToolUsageEnabled()) {
            return List.of();
        }

        try {
            ToolManagerClient.GetToolsResponse response = toolManagerClient
                .getTools(
                    agent.getToolManagerUrl(),
                    session.getSessionId(),
                    agent.getAgentId(),
                    Map.of("enabled", true)
                )
                .block();

            if (response != null && response.getTools() != null) {
                return response.getTools().stream()
                    .map(tool -> Map.<String, Object>of(
                        "name", tool.getName(),
                        "description", tool.getDescription(),
                        "parameters", tool.getParameters()
                    ))
                    .toList();
            }

        } catch (Exception e) {
            log.error("Failed to fetch tools", e);
        }

        return List.of();
    }
}
